/*
 * Copyright (C) 2007 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.android.dx.dex.code;

import java.util.ArrayList;
import java.util.BitSet;
import java.util.HashSet;

import com.android.dx.dex.DexOptions;
import com.android.dx.io.opcode.Opcodes;
import com.android.dx.rop.code.LocalItem;
import com.android.dx.rop.code.RegisterSpec;
import com.android.dx.rop.code.RegisterSpecList;
import com.android.dx.rop.code.RegisterSpecSet;
import com.android.dx.rop.code.SourcePosition;
import com.android.dx.rop.cst.Constant;
import com.android.dx.rop.cst.CstMemberRef;
import com.android.dx.rop.cst.CstString;
import com.android.dx.rop.cst.CstType;
import com.android.dx.rop.type.Type;
import com.android.dx.util.DexException;

/**
 * Processor for instruction lists, which takes a "first cut" of instruction
 * selection as a basis and produces a "final cut" in the form of a
 * {@link DalvInsnList} instance.
 */
public final class OutputFinisher {

	/** {@code non-null;} options for dex output */
	private final DexOptions dexOptions;

	/**
	 * {@code >= 0;} register count for the method, not including any extra
	 * "reserved" registers needed to translate "difficult" instructions
	 */
	private final int unreservedRegCount;

	/** {@code non-null;} the list of instructions, per se */
	private ArrayList<DalvInsn> insns;

	/** whether any instruction has position info */
	private boolean hasAnyPositionInfo;

	/** whether any instruction has local variable info */
	private boolean hasAnyLocalInfo;

	/**
	 * {@code >= 0;} the count of reserved registers (low-numbered registers
	 * used when expanding instructions that can't be represented simply);
	 * becomes valid after a call to {@link #massageInstructions}
	 */
	private int reservedCount;

	/**
	 * Constructs an instance. It initially contains no instructions.
	 * 
	 * @param dexOptions
	 *            {@code non-null;} options for dex output
	 * @param regCount
	 *            {@code >= 0;} register count for the method
	 * @param initialCapacity
	 *            {@code >= 0;} initial capacity of the instructions list
	 */
	public OutputFinisher(DexOptions dexOptions, int initialCapacity,
			int regCount) {
		this.dexOptions = dexOptions;
		this.unreservedRegCount = regCount;
		this.insns = new ArrayList<DalvInsn>(initialCapacity);
		this.reservedCount = -1;
		this.hasAnyPositionInfo = false;
		this.hasAnyLocalInfo = false;
	}

	/**
	 * Returns whether any of the instructions added to this instance come with
	 * position info.
	 * 
	 * @return whether any of the instructions added to this instance come with
	 *         position info
	 */
	public boolean hasAnyPositionInfo() {
		return hasAnyPositionInfo;
	}

	/**
	 * Returns whether this instance has any local variable information.
	 * 
	 * @return whether this instance has any local variable information
	 */
	public boolean hasAnyLocalInfo() {
		return hasAnyLocalInfo;
	}

	/**
	 * Helper for {@link #add} which scrutinizes a single instruction for local
	 * variable information.
	 * 
	 * @param insn
	 *            {@code non-null;} instruction to scrutinize
	 * @return {@code true} iff the instruction refers to any named locals
	 */
	private static boolean hasLocalInfo(DalvInsn insn) {
		if (insn instanceof LocalSnapshot) {
			RegisterSpecSet specs = ((LocalSnapshot) insn).getLocals();
			int size = specs.size();
			for (int i = 0; i < size; i++) {
				if (hasLocalInfo(specs.get(i))) {
					return true;
				}
			}
		} else if (insn instanceof LocalStart) {
			RegisterSpec spec = ((LocalStart) insn).getLocal();
			if (hasLocalInfo(spec)) {
				return true;
			}
		}

		return false;
	}

	/**
	 * Helper for {@link #hasAnyLocalInfo} which scrutinizes a single register
	 * spec.
	 * 
	 * @param spec
	 *            {@code non-null;} spec to scrutinize
	 * @return {@code true} iff the spec refers to any named locals
	 */
	private static boolean hasLocalInfo(RegisterSpec spec) {
		return (spec != null) && (spec.getLocalItem().getName() != null);
	}

	/**
	 * Returns the set of all constants referred to by instructions added to
	 * this instance.
	 * 
	 * @return {@code non-null;} the set of constants
	 */
	public HashSet<Constant> getAllConstants() {
		HashSet<Constant> result = new HashSet<Constant>(20);

		for (DalvInsn insn : insns) {
			addConstants(result, insn);
		}

		return result;
	}

	/**
	 * Helper for {@link #getAllConstants} which adds all the info for a single
	 * instruction.
	 * 
	 * @param result
	 *            {@code non-null;} result set to add to
	 * @param insn
	 *            {@code non-null;} instruction to scrutinize
	 */
	private static void addConstants(HashSet<Constant> result, DalvInsn insn) {
		if (insn instanceof CstInsn) {
			Constant cst = ((CstInsn) insn).getConstant();
			result.add(cst);
		} else if (insn instanceof LocalSnapshot) {
			RegisterSpecSet specs = ((LocalSnapshot) insn).getLocals();
			int size = specs.size();
			for (int i = 0; i < size; i++) {
				addConstants(result, specs.get(i));
			}
		} else if (insn instanceof LocalStart) {
			RegisterSpec spec = ((LocalStart) insn).getLocal();
			addConstants(result, spec);
		}
	}

	/**
	 * Helper for {@link #getAllConstants} which adds all the info for a single
	 * {@code RegisterSpec}.
	 * 
	 * @param result
	 *            {@code non-null;} result set to add to
	 * @param spec
	 *            {@code null-ok;} register spec to add
	 */
	private static void addConstants(HashSet<Constant> result, RegisterSpec spec) {
		if (spec == null) {
			return;
		}

		LocalItem local = spec.getLocalItem();
		CstString name = local.getName();
		CstString signature = local.getSignature();
		Type type = spec.getType();

		if (type != Type.KNOWN_NULL) {
			result.add(CstType.intern(type));
		}

		if (name != null) {
			result.add(name);
		}

		if (signature != null) {
			result.add(signature);
		}
	}

	/**
	 * Adds an instruction to the output.
	 * 
	 * @param insn
	 *            {@code non-null;} the instruction to add
	 */
	public void add(DalvInsn insn) {
		insns.add(insn);
		updateInfo(insn);
	}

	/**
	 * Inserts an instruction in the output at the given offset.
	 * 
	 * @param at
	 *            {@code >= 0;} what index to insert at
	 * @param insn
	 *            {@code non-null;} the instruction to insert
	 */
	public void insert(int at, DalvInsn insn) {
		insns.add(at, insn);
		updateInfo(insn);
	}

	/**
	 * Helper for {@link #add} and {@link #insert}, which updates the position
	 * and local info flags.
	 * 
	 * @param insn
	 *            {@code non-null;} an instruction that was just introduced
	 */
	private void updateInfo(DalvInsn insn) {
		if (!hasAnyPositionInfo) {
			SourcePosition pos = insn.getPosition();
			if (pos.getLine() >= 0) {
				hasAnyPositionInfo = true;
			}
		}

		if (!hasAnyLocalInfo) {
			if (hasLocalInfo(insn)) {
				hasAnyLocalInfo = true;
			}
		}
	}

	/**
	 * Reverses a branch which is buried a given number of instructions backward
	 * in the output. It is illegal to call this unless the indicated
	 * instruction really is a reversible branch.
	 * 
	 * @param which
	 *            how many instructions back to find the branch; {@code 0} is
	 *            the most recently added instruction, {@code 1} is the
	 *            instruction before that, etc.
	 * @param newTarget
	 *            {@code non-null;} the new target for the reversed branch
	 */
	public void reverseBranch(int which, CodeAddress newTarget) {
		int size = insns.size();
		int index = size - which - 1;
		TargetInsn targetInsn;

		try {
			targetInsn = (TargetInsn) insns.get(index);
		} catch (IndexOutOfBoundsException ex) {
			// Translate the exception.
			throw new IllegalArgumentException("too few instructions");
		} catch (ClassCastException ex) {
			// Translate the exception.
			throw new IllegalArgumentException("non-reversible instruction");
		}

		/*
		 * No need to call this.set(), since the format and other info are the
		 * same.
		 */
		insns.set(index, targetInsn.withNewTargetAndReversed(newTarget));
	}

	/**
	 * Assigns indices in all instructions that need them, using the given
	 * callback to perform lookups. This should be called before calling
	 * {@link #finishProcessingAndGetList}.
	 * 
	 * @param callback
	 *            {@code non-null;} callback object
	 */
	public void assignIndices(DalvCode.AssignIndicesCallback callback) {
		for (DalvInsn insn : insns) {
			if (insn instanceof CstInsn) {
				assignIndices((CstInsn) insn, callback);
			}
		}
	}

	/**
	 * Helper for {@link #assignIndices} which does assignment for one
	 * instruction.
	 * 
	 * @param insn
	 *            {@code non-null;} the instruction
	 * @param callback
	 *            {@code non-null;} the callback
	 */
	private static void assignIndices(CstInsn insn,
			DalvCode.AssignIndicesCallback callback) {
		Constant cst = insn.getConstant();
		int index = callback.getIndex(cst);

		if (index >= 0) {
			insn.setIndex(index);
		}

		if (cst instanceof CstMemberRef) {
			CstMemberRef member = (CstMemberRef) cst;
			CstType definer = member.getDefiningClass();
			index = callback.getIndex(definer);
			if (index >= 0) {
				insn.setClassIndex(index);
			}
		}
	}

	/**
	 * Does final processing on this instance and gets the output as a
	 * {@link DalvInsnList}. Final processing consists of:
	 * <ul>
	 * <li>optionally renumbering registers (to make room as needed for expanded
	 * instructions)</li>
	 * <li>picking a final opcode for each instruction</li>
	 * <li>rewriting instructions, because of register number, constant pool
	 * index, or branch target size issues</li>
	 * <li>assigning final addresses</li>
	 * </ul>
	 * <p>
	 * <b>Note:</b> This method may only be called once per instance of this
	 * class.
	 * </p>
	 * 
	 * @return {@code non-null;} the output list
	 * @throws UnsupportedOperationException
	 *             if this method has already been called
	 */
	public DalvInsnList finishProcessingAndGetList() {
		if (reservedCount >= 0) {
			throw new UnsupportedOperationException("already processed");
		}

		Dop[] opcodes = makeOpcodesArray();
		reserveRegisters(opcodes);
		massageInstructions(opcodes);
		assignAddressesAndFixBranches();

		return DalvInsnList.makeImmutable(insns, reservedCount
				+ unreservedRegCount);
	}

	/**
	 * Helper for {@link #finishProcessingAndGetList}, which extracts the opcode
	 * out of each instruction into a separate array, to be further manipulated
	 * as things progress.
	 * 
	 * @return {@code non-null;} the array of opcodes
	 */
	private Dop[] makeOpcodesArray() {
		int size = insns.size();
		Dop[] result = new Dop[size];

		for (int i = 0; i < size; i++) {
			result[i] = insns.get(i).getOpcode();
		}

		return result;
	}

	/**
	 * Helper for {@link #finishProcessingAndGetList}, which figures out how
	 * many reserved registers are required and then reserving them. It also
	 * updates the given {@code opcodes} array so as to avoid extra work when
	 * constructing the massaged instruction list.
	 * 
	 * @param opcodes
	 *            {@code non-null;} array of per-instruction opcode selections
	 */
	private void reserveRegisters(Dop[] opcodes) {
		int oldReservedCount = (reservedCount < 0) ? 0 : reservedCount;

		/*
		 * Call calculateReservedCount() and then perform register reservation,
		 * repeatedly until no new reservations happen.
		 */
		for (;;) {
			int newReservedCount = calculateReservedCount(opcodes);
			if (oldReservedCount >= newReservedCount) {
				break;
			}

			int reservedDifference = newReservedCount - oldReservedCount;
			int size = insns.size();

			for (int i = 0; i < size; i++) {
				/*
				 * CodeAddress instance identity is used to link TargetInsns to
				 * their targets, so it is inappropriate to make replacements,
				 * and they don't have registers in any case. Hence, the
				 * instanceof test below.
				 */
				DalvInsn insn = insns.get(i);
				if (!(insn instanceof CodeAddress)) {
					/*
					 * No need to call this.set() since the format and other
					 * info are the same.
					 */
					insns.set(i, insn.withRegisterOffset(reservedDifference));
				}
			}

			oldReservedCount = newReservedCount;
		}

		reservedCount = oldReservedCount;
	}

	/**
	 * Helper for {@link #reserveRegisters}, which does one pass over the
	 * instructions, calculating the number of registers that need to be
	 * reserved. It also updates the {@code opcodes} list to help avoid extra
	 * work in future register reservation passes.
	 * 
	 * @param opcodes
	 *            {@code non-null;} array of per-instruction opcode selections
	 * @return {@code >= 0;} the count of reserved registers
	 */
	private int calculateReservedCount(Dop[] opcodes) {
		int size = insns.size();

		/*
		 * Potential new value of reservedCount, which gets updated in the
		 * following loop. It starts out with the existing reservedCount and
		 * gets increased if it turns out that additional registers need to be
		 * reserved.
		 */
		int newReservedCount = reservedCount;

		for (int i = 0; i < size; i++) {
			DalvInsn insn = insns.get(i);
			Dop originalOpcode = opcodes[i];
			Dop newOpcode = findOpcodeForInsn(insn, originalOpcode);

			if (newOpcode == null) {
				/*
				 * The instruction will need to be expanded, so find the
				 * expanded opcode and reserve registers for it.
				 */
				Dop expandedOp = findExpandedOpcodeForInsn(insn);
				BitSet compatRegs = expandedOp.getFormat().compatibleRegs(insn);
				int reserve = insn.getMinimumRegisterRequirement(compatRegs);
				if (reserve > newReservedCount) {
					newReservedCount = reserve;
				}
			} else if (originalOpcode == newOpcode) {
				continue;
			}

			opcodes[i] = newOpcode;
		}

		return newReservedCount;
	}

	/**
	 * Attempts to fit the given instruction into a specific opcode, returning
	 * the opcode whose format that the instruction fits into or {@code null} to
	 * indicate that the instruction will need to be expanded. This fitting
	 * process starts with the given opcode as a first "best guess" and then
	 * pessimizes from there if necessary.
	 * 
	 * @param insn
	 *            {@code non-null;} the instruction in question
	 * @param guess
	 *            {@code null-ok;} the current guess as to the best opcode;
	 *            {@code null} means that no simple opcode fits
	 * @return {@code null-ok;} a possibly-different opcode; either a
	 *         {@code non-null} good fit or {@code null} to indicate that no
	 *         simple opcode fits
	 */
	private Dop findOpcodeForInsn(DalvInsn insn, Dop guess) {
		/*
		 * Note: The initial guess might be null, meaning that an earlier call
		 * to this method already determined that there was no possible simple
		 * opcode fit.
		 */

		while (guess != null) {
			if (guess.getFormat().isCompatible(insn)) {
				break;
			}

			guess = Dops.getNextOrNull(guess, dexOptions);
		}

		return guess;
	}

	/**
	 * Finds the proper opcode for the given instruction, ignoring register
	 * constraints.
	 * 
	 * @param insn
	 *            {@code non-null;} the instruction in question
	 * @return {@code non-null;} the opcode that fits
	 */
	private Dop findExpandedOpcodeForInsn(DalvInsn insn) {
		Dop result = findOpcodeForInsn(insn.getLowRegVersion(),
				insn.getOpcode());
		if (result == null) {
			throw new DexException("No expanded opcode for " + insn);
		}
		return result;
	}

	/**
	 * Helper for {@link #finishProcessingAndGetList}, which goes through each
	 * instruction in the output, making sure its opcode can accomodate its
	 * arguments. In cases where the opcode is unable to do so, this replaces
	 * the instruction with a larger instruction with identical semantics that
	 * <i>will</i> work.
	 * <p>
	 * This method may also reserve a number of low-numbered registers,
	 * renumbering the instructions' original registers, in order to have
	 * register space available in which to move very-high registers when
	 * expanding instructions into multi-instruction sequences. This expansion
	 * is done when no simple instruction format can be found for a given
	 * instruction that is able to accomodate that instruction's registers.
	 * </p>
	 * <p>
	 * This method ignores issues of branch target size, since final addresses
	 * aren't known at the point that this method is called.
	 * </p>
	 * 
	 * @param opcodes
	 *            {@code non-null;} array of per-instruction opcode selections
	 */
	private void massageInstructions(Dop[] opcodes) {
		if (reservedCount == 0) {
			/*
			 * The easy common case: No registers were reserved, so we merely
			 * need to replace any instructions whose format (and hence whose
			 * opcode) changed during the reservation pass, but all instructions
			 * will stay at their original indices, and the instruction list
			 * doesn't grow.
			 */
			int size = insns.size();

			for (int i = 0; i < size; i++) {
				DalvInsn insn = insns.get(i);
				Dop originalOpcode = insn.getOpcode();
				Dop currentOpcode = opcodes[i];

				if (originalOpcode != currentOpcode) {
					insns.set(i, insn.withOpcode(currentOpcode));
				}
			}
		} else {
			/*
			 * The difficult uncommon case: Some instructions have to be
			 * expanded to deal with high registers.
			 */
			insns = performExpansion(opcodes);
		}
	}

	/**
	 * Helper for {@link #massageInstructions}, which constructs a replacement
	 * list, where each {link DalvInsn} instance that couldn't be represented
	 * simply (due to register representation problems) is expanded into a
	 * series of instances that together perform the proper function.
	 * 
	 * @param opcodes
	 *            {@code non-null;} array of per-instruction opcode selections
	 * @return {@code non-null;} the replacement list
	 */
	private ArrayList<DalvInsn> performExpansion(Dop[] opcodes) {
		int size = insns.size();
		ArrayList<DalvInsn> result = new ArrayList<DalvInsn>(size * 2);

		for (int i = 0; i < size; i++) {
			DalvInsn insn = insns.get(i);
			Dop originalOpcode = insn.getOpcode();
			Dop currentOpcode = opcodes[i];
			DalvInsn prefix;
			DalvInsn suffix;

			if (currentOpcode != null) {
				// No expansion is necessary.
				prefix = null;
				suffix = null;
			} else {
				// Expansion is required.
				currentOpcode = findExpandedOpcodeForInsn(insn);
				BitSet compatRegs = currentOpcode.getFormat().compatibleRegs(
						insn);
				prefix = insn.expandedPrefix(compatRegs);
				suffix = insn.expandedSuffix(compatRegs);

				// Expand necessary registers to fit the new format
				insn = insn.expandedVersion(compatRegs);
			}

			if (prefix != null) {
				result.add(prefix);
			}

			if (currentOpcode != originalOpcode) {
				insn = insn.withOpcode(currentOpcode);
			}
			result.add(insn);

			if (suffix != null) {
				result.add(suffix);
			}
		}

		return result;
	}

	/**
	 * Helper for {@link #finishProcessingAndGetList}, which assigns addresses
	 * to each instruction, possibly rewriting branches to fix ones that
	 * wouldn't otherwise be able to reach their targets.
	 */
	private void assignAddressesAndFixBranches() {
		for (;;) {
			assignAddresses();
			if (!fixBranches()) {
				break;
			}
		}
	}

	/**
	 * Helper for {@link #assignAddressesAndFixBranches}, which assigns an
	 * address to each instruction, in order.
	 */
	private void assignAddresses() {
		int address = 0;
		int size = insns.size();

		for (int i = 0; i < size; i++) {
			DalvInsn insn = insns.get(i);
			insn.setAddress(address);
			address += insn.codeSize();
		}
	}

	/**
	 * Helper for {@link #assignAddressesAndFixBranches}, which checks the
	 * branch target size requirement of each branch instruction to make sure it
	 * fits. For instructions that don't fit, this rewrites them to use a
	 * {@code goto} of some sort. In the case of a conditional branch that
	 * doesn't fit, the sense of the test is reversed in order to branch around
	 * a {@code goto} to the original target.
	 * 
	 * @return whether any branches had to be fixed
	 */
	private boolean fixBranches() {
		int size = insns.size();
		boolean anyFixed = false;

		for (int i = 0; i < size; i++) {
			DalvInsn insn = insns.get(i);
			if (!(insn instanceof TargetInsn)) {
				// This loop only needs to inspect TargetInsns.
				continue;
			}

			Dop opcode = insn.getOpcode();
			TargetInsn target = (TargetInsn) insn;

			if (opcode.getFormat().branchFits(target)) {
				continue;
			}

			if (opcode.getFamily() == Opcodes.GOTO) {
				// It is a goto; widen it if possible.
				opcode = findOpcodeForInsn(insn, opcode);
				if (opcode == null) {
					/*
					 * The branch is already maximally large. This should only
					 * be possible if a method somehow manages to have more than
					 * 2^31 code units.
					 */
					throw new UnsupportedOperationException("method too long");
				}
				insns.set(i, insn.withOpcode(opcode));
			} else {
				/*
				 * It is a conditional: Reverse its sense, and arrange for it to
				 * branch around an absolute goto to the original branch target.
				 * 
				 * Note: An invariant of the list being processed is that every
				 * TargetInsn is followed by a CodeAddress. Hence, it is always
				 * safe to get the next element after a TargetInsn and cast it
				 * to CodeAddress, as is happening a few lines down.
				 * 
				 * Also note: Size gets incremented by one here, as we have --
				 * in the net -- added one additional element to the list, so we
				 * increment i to match. The added and changed elements will be
				 * inspected by a repeat call to this method after this
				 * invocation returns.
				 */
				CodeAddress newTarget;
				try {
					newTarget = (CodeAddress) insns.get(i + 1);
				} catch (IndexOutOfBoundsException ex) {
					// The TargetInsn / CodeAddress invariant was violated.
					throw new IllegalStateException(
							"unpaired TargetInsn (dangling)");
				} catch (ClassCastException ex) {
					// The TargetInsn / CodeAddress invariant was violated.
					throw new IllegalStateException("unpaired TargetInsn");
				}
				TargetInsn gotoInsn = new TargetInsn(Dops.GOTO,
						target.getPosition(), RegisterSpecList.EMPTY,
						target.getTarget());
				insns.set(i, gotoInsn);
				insns.add(i, target.withNewTargetAndReversed(newTarget));
				size++;
				i++;
			}

			anyFixed = true;
		}

		return anyFixed;
	}
}
